package com.njupt.wuaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.njupt.wuaiagent.agent.model.AgentState;
import dev.langchain4j.agent.tool.P;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @Author: wujiaming
 * @CreateTime: 2025/5/22 19:55
 * @Description: 负责实现工具调用能力，继承自ReActAgent，具体实现了think 和act两个抽象方法。
 * 可以用作创建实例的父类
 * @Version: 1.0
 */


@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent{

    //可用的工具
    private final ToolCallback[] availableTools;

    //保存工具调用的响应信息（要调用哪些工具）
    private ChatResponse toolCallChatResponse;

    //工具调用的管理者（Spring AI用ToolCallingManager来管理要调用的工具，工具调用管理者）
    private final ToolCallingManager toolCallingManager;

    //静止内置的工具调用机制，自己维护上下文
    //SpringAI实现自主控制工具调用：构造一个对话选项，在对话选项中关闭自动工具调用，然后把对话选项传递给prompt
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools){
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();

        //禁止Spring AI 内置的工具调用机制，自己维护选项和消息上下文
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
    }

    /**
     * 处理当前的状态的状态决定下一步行动
     *
     * @return 是否需要执行行动
     */
    @Override
    public boolean think() {
        //1、获取最新的对话上下文记忆。校验提示词，上一步的提示词拼接到对话上下文中
        if(StrUtil.isNotBlank(getNextStepPrompt())){
            UserMessage userMessage = new UserMessage(getNextStepPrompt());
            getMessageList().add(userMessage);

        }

        //2、调用 AI大模型，获取大模型到调用的工具列表
        List<Message> messageList = getMessageList();
        Prompt prompt = new Prompt(getMessageList(),chatOptions);
        try {
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();


            // 记录响应，用于等下的Act
            this.toolCallChatResponse = chatResponse;
            //3、解析调用的结果，获取要调用的工具
            //助手信息，相当于AI的思考结果，AI中间返回的结果
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            //获取AI想要调用的工具
            List<AssistantMessage.ToolCall> toolCalls = assistantMessage.getToolCalls();

            //输出信息
            String result = assistantMessage.getText();
            log.info(getName() + "的思考：" + result);
            log.info(getName() + "选择了 ：" + toolCalls.size()  + "个工具来使用");
            String toolCallInfo = toolCalls.stream()
                    .map(toolCall -> {
                        return String.format("工具的名称： %s，参数 %s", toolCall.name(), toolCall.arguments());
                    }).collect(Collectors.joining("\n"));
            log.info(toolCallInfo);

            //判断是否要调用工具
            //不需要调用工具，返回false
            if(toolCallInfo.isEmpty()){
                //将AI的回复以及工具调用的结果放在记忆上下文信息中。只有不调用工具时，才需要手动记录助手消息
                getMessageList().add(assistantMessage);
                return false;
            }else{
                // 需要调用工具时，无需记录助手消息，因为调用工具时会自动记录
                return true;
            }
        } catch (Exception e) {
            log.error(getName() + "的思考过程遇到了问题：" + e.getMessage());
            getMessageList().add(new AssistantMessage("处理时遇到了错误:" + e.getMessage()));
            return false;
        }

    }

    /**
     * 执行工具调用并且处理结果
     *
     * @return 执行结果
     */
    @Override
    public String act() {

        if (!toolCallChatResponse.hasToolCalls()) {
            return "没有工具需要调用";
        }
        //调用工具
        Prompt prompt = new Prompt(getMessageList(),this.chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        //记录消息上下文，conversationHistory 已经包含了助手消息和工具调用返回的结果
        setMessageList(toolExecutionResult.conversationHistory());
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());

        //判断是否调用了终止工具
        boolean terminateToolCalled = toolResponseMessage.getResponses().stream()
                .anyMatch(toolResponse -> toolResponse.name().equals("doTerminate"));
        String results = toolResponseMessage.getResponses().stream()
                .map(toolResponse -> {
                    return "工具" + toolResponse.name() + " 返回结果：" + toolResponse.responseData();
                }).collect(Collectors.joining("\n"));
        if (terminateToolCalled) {

            //任务结束，更改状态
            setState(AgentState.FINISHED);
        }

        log.info(results);
        return results;
    }
}
